#! /usr/bin/env python
# -*- coding:UTF-8 -*-
# 上面这两行代码必须要加上，否则在ros当中无法直接运行
# import rospy

# 8.25
# 红绿灯代码测试
import cv2
import numpy as np


# 这个是你自己写的红绿灯检测函数

# 这里定义红绿等的检测函数
def light_detection(involved_light_img):
    # -------------------------------------------------------------------  # 参数的设定
    # 判定标志(初始化为None)
    Light_R_flag = None
    # 从参数当中读取到图片
    image_BGR = involved_light_img
    # 将图片转换为HSV颜色空间
    image_HSV = cv2.cvtColor(src=image_BGR, code=cv2.COLOR_BGR2HSV)

    # 定义红色的HSV范围
    # 由于红色是在H:0和H:180附近，因此会特殊一点，这里创建双向的阈值
    Red_lower_pos = np.array([0, 130, 60])
    Red_upper_pos = np.array([5, 255, 255])
    Red_lower_neg = np.array([170, 130, 60])
    Red_upper_neg = np.array([180, 255, 255])

    # 定义绿色HSV的范围
    Green_lower = np.array([55, 85, 60])
    Green_upper = np.array([75, 255, 255])

    # -----------------------------------------------------------------------   # 首先是红色的检测
    # 这里创建两个红色的掩膜，使用两个阈值
    Red_mask_pos = cv2.inRange(src=image_HSV,
                               lowerb=Red_lower_pos,
                               upperb=Red_upper_pos)
    Red_mask_neg = cv2.inRange(src=image_HSV,
                               lowerb=Red_lower_neg,
                               upperb=Red_upper_neg)
    # 将红色掩膜相加
    Red_mask = cv2.add(Red_mask_pos, Red_mask_neg)

    # ------------------------------------------------------
    # 这里你可以选择去不去白色区域的黑色噪点
    # # 这里生成用于形态学操作的卷积核
    # kernel_set = cv2.getStructuringElement(shape=cv2.MORPH_RECT,
    #                                        ksize=(7, 7),
    #                                        anchor=(-1, -1)
    #                                        )
    # # 这里进行形态学闭操作
    # Red_mask = cv2.morphologyEx(src=Red_mask,
    #                             op=cv2.MORPH_CLOSE,
    #                             kernel=kernel_set,
    #                             iterations=3)
    # ---------------------------------------------------------------

    # 这里计算红色掩膜当中有多少白点,有两种方式，如下：
    # 方式1
    # R_flag = Red_mask[:, :] == 255
    # R_Count = len(Red_mask[R_flag])
    # print("count_R:", R_Count)

    # 方式2
    R_Count = cv2.countNonZero(Red_mask)
    # print("count_R:", R_Count)

    # -----------------------------------------------------------------------    # 往下是绿色的检测

    # 创建绿色的掩膜
    Green_mask = cv2.inRange(src=image_HSV,
                             lowerb=Green_lower,
                             upperb=Green_upper)

    # 这里计算绿色掩膜当中有多少白点
    # 方式1
    # G_flag = Green_mask[:, :] == 255
    # G_Count = len(Green_mask[G_flag])
    # print("count_G:", G_Count)

    # 方式2
    G_Count = cv2.countNonZero(Green_mask)
    # print("count_G:", G_Count)

    # ---------------------------------------------       # 掩膜的显示

    # 显示红色的掩膜
    # cv2.imshow('mask_red', Red_mask)
    # 显示绿色的掩膜
    # cv2.imshow('mask_green', Green_mask)

    # --------------------------------------------         # 红绿灯的判定

    # 接下来是红绿灯判定
    # 这里设定的阈值就是2000
    if R_Count >= 2000 and R_Count > G_Count:
        print("Red: ", R_Count)
        return 0
    elif G_Count >= 2000 and G_Count > R_Count:
        print("Green: ", G_Count)
        return 1
    else:
        print("None")
        return 0


# 开头的路线选择
def RG_light_choice(involved_light_img):
    # -------------------------------------------------------------------  # 参数的设定
    # 判定标志(初始化为None)
    Light_R_flag = None
    # 从参数当中读取到图片
    image_BGR = involved_light_img

    # 将图片转换为HSV颜色空间
    image_HSV = cv2.cvtColor(src=image_BGR, code=cv2.COLOR_BGR2HSV)

    # 拿到左边的图
    img_left_HSV = image_HSV[:, 0:320]
    # 拿到右边的图像
    img_right_HSV = image_HSV[:, 320:640]

    # 定义红色的HSV范围
    # 由于红色是在H:0和H:180附近，因此会特殊一点，这里创建双向的阈值
    Red_lower_pos = np.array([0, 130, 60])
    Red_upper_pos = np.array([5, 255, 255])
    Red_lower_neg = np.array([170, 130, 60])
    Red_upper_neg = np.array([180, 255, 255])

    # 定义绿色HSV的范围
    Green_lower = np.array([55, 85, 60])
    Green_upper = np.array([75, 255, 255])

    # -----------------------------------------------------------------------   # 首先是红色的检测
    # 这里创建两个红色的掩膜，使用两个阈值
    Red_mask_pos = cv2.inRange(src=image_HSV,
                               lowerb=Red_lower_pos,
                               upperb=Red_upper_pos)
    Red_mask_neg = cv2.inRange(src=image_HSV,
                               lowerb=Red_lower_neg,
                               upperb=Red_upper_neg)
    # 将红色掩膜相加
    Red_mask = cv2.add(Red_mask_pos, Red_mask_neg)

    # ------------------------------------------------------
    # 这里你可以选择去不去白色区域的黑色噪点
    # # 这里生成用于形态学操作的卷积核
    # kernel_set = cv2.getStructuringElement(shape=cv2.MORPH_RECT,
    #                                        ksize=(7, 7),
    #                                        anchor=(-1, -1)
    #                                        )
    # # 这里进行形态学闭操作
    # Red_mask = cv2.morphologyEx(src=Red_mask,
    #                             op=cv2.MORPH_CLOSE,
    #                             kernel=kernel_set,
    #                             iterations=3)
    # ---------------------------------------------------------------

    # 这里计算红色掩膜当中有多少白点,有两种方式，如下：
    # 方式1
    # R_flag = Red_mask[:, :] == 255
    # R_Count = len(Red_mask[R_flag])
    # print("count_R:", R_Count)

    # 方式2
    R_Count = cv2.countNonZero(Red_mask)
    # print("count_R:", R_Count)

    # -----------------------------------------------------------------------    # 往下是绿色的检测

    # 创建绿色的掩膜
    Green_mask = cv2.inRange(src=image_HSV,
                             lowerb=Green_lower,
                             upperb=Green_upper)

    Green_mask_left = cv2.inRange(src=img_left_HSV,
                                  lowerb=Green_lower,
                                  upperb=Green_upper)

    Green_mask_right = cv2.inRange(src=img_right_HSV,
                                   lowerb=Green_lower,
                                   upperb=Green_upper)

    # 这里计算绿色掩膜当中有多少白点
    # 方式1
    # G_flag = Green_mask[:, :] == 255
    # G_Count = len(Green_mask[G_flag])
    # print("count_G:", G_Count)

    # 方式2
    G_Count = cv2.countNonZero(Green_mask)

    G_count_left = cv2.countNonZero(Green_mask_left)

    G_count_right = cv2.countNonZero(Green_mask_right)

    # print("count_G:", G_Count)

    # ---------------------------------------------       # 掩膜的显示

    # # 显示红色的掩膜
    # cv2.imshow('mask_red', Red_mask)
    # # 显示绿色的掩膜
    # cv2.imshow('mask_green', Green_mask)

    # --------------------------------------------         # 红绿灯的判定

    # 接下来是红绿灯判定
    # 这里设定的阈值就是2000
    if R_Count >= 2000 and R_Count > G_Count:
        print("Red: ", R_Count)
        return 0
    elif G_Count >= 2000 and G_Count > R_Count:

        if G_count_left > G_count_right:
            print('1111111')
            return 1

        elif G_count_left < G_count_right:
            print('2222222')
            return 2
        else:
            print('0000000')
            return 0
    else:
        print("None")
        return 0


# 中间的停止路线选择
def RG_light_choice_during_process(img):
    # -------------------------------------------------------------------  # 参数的设定
    # 从参数当中读取到图片
    image_BGR = img

    # 将图片转换为HSV颜色空间
    image_HSV = cv2.cvtColor(src=image_BGR, code=cv2.COLOR_BGR2HSV)

    # 拿到左边的图
    img_left_HSV = image_HSV[:, 0:320]
    # 拿到右边的图像
    img_right_HSV = image_HSV[:, 320:640]

    # 定义红色的HSV范围
    # 由于红色是在H:0和H:180附近，因此会特殊一点，这里创建双向的阈值
    Red_lower_pos = np.array([0, 130, 60])
    Red_upper_pos = np.array([5, 255, 255])
    Red_lower_neg = np.array([170, 130, 60])
    Red_upper_neg = np.array([180, 255, 255])

    # 定义绿色HSV的范围
    Green_lower = np.array([55, 85, 60])
    Green_upper = np.array([75, 255, 255])

    # -----------------------------------------------------------------------   # 首先是红色的检测
    # 这里创建两个红色的掩膜，使用两个阈值
    Red_mask_pos = cv2.inRange(src=image_HSV,
                               lowerb=Red_lower_pos,
                               upperb=Red_upper_pos)
    Red_mask_neg = cv2.inRange(src=image_HSV,
                               lowerb=Red_lower_neg,
                               upperb=Red_upper_neg)
    # 将红色掩膜相加
    Red_mask = cv2.add(Red_mask_pos, Red_mask_neg)

    # ------------------------------------------------------
    # 这里你可以选择去不去白色区域的黑色噪点
    # # 这里生成用于形态学操作的卷积核
    # kernel_set = cv2.getStructuringElement(shape=cv2.MORPH_RECT,
    #                                        ksize=(7, 7),
    #                                        anchor=(-1, -1)
    #                                        )
    # # 这里进行形态学闭操作
    # Red_mask = cv2.morphologyEx(src=Red_mask,
    #                             op=cv2.MORPH_CLOSE,
    #                             kernel=kernel_set,
    #                             iterations=3)
    # ---------------------------------------------------------------

    # 这里计算红色掩膜当中有多少白点,有两种方式，如下：
    # 方式1
    # R_flag = Red_mask[:, :] == 255
    # R_Count = len(Red_mask[R_flag])
    # print("count_R:", R_Count)

    # 方式2
    R_Count = cv2.countNonZero(Red_mask)
    # print("count_R:", R_Count)

    # -----------------------------------------------------------------------    # 往下是绿色的检测

    # 创建绿色的掩膜
    Green_mask = cv2.inRange(src=image_HSV,
                             lowerb=Green_lower,
                             upperb=Green_upper)

    Green_mask_left = cv2.inRange(src=img_left_HSV,
                                  lowerb=Green_lower,
                                  upperb=Green_upper)

    Green_mask_right = cv2.inRange(src=img_right_HSV,
                                   lowerb=Green_lower,
                                   upperb=Green_upper)

    # 这里计算绿色掩膜当中有多少白点
    # 方式1
    # G_flag = Green_mask[:, :] == 255
    # G_Count = len(Green_mask[G_flag])
    # print("count_G:", G_Count)

    # 方式2
    G_Count = cv2.countNonZero(Green_mask)

    G_count_left = cv2.countNonZero(Green_mask_left)

    G_count_right = cv2.countNonZero(Green_mask_right)

    # print("count_G:", G_Count)

    # ---------------------------------------------       # 掩膜的显示

    # # 显示红色的掩膜
    # cv2.imshow('mask_red', Red_mask)
    # # 显示绿色的掩膜
    # cv2.imshow('mask_green', Green_mask)

    # --------------------------------------------         # 红绿灯的判定

    # 接下来是红绿灯判定
    # 这里设定的阈值就是2000
    if R_Count >= 2000 and R_Count > G_Count:
        print("Red: ", R_Count)
        return 444
    elif G_Count >= 2000 and G_Count > R_Count:

        if G_count_left > G_count_right:
            print('1111111')
            return 1

        elif G_count_left < G_count_right:
            print('2222222')
            return 2
        else:
            print('0000000')
            return 0
    else:
        print("None")
        return 0


    pass


# 这里定义摄像头畸变矫正函数
def Straighten_the_camera(original_img):
    # 这里对摄像头的视频帧进行矫正
    # 定义相机内参矩阵，畸变矩阵，用来矫正相机畸变：相机型号HF867
    CamMax = np.array([[367.87792233, 0, 320.00047236],
                       [0, 370.36759206, 240.00000000],
                       [0, 0, 1]])

    disCef = np.array([-0.3369, 0.0870, 0, 0, 0])

    # 这里就是矫正摄像头畸变的操作
    global und_st
    und_st = cv2.undistort(src=original_img,
                           cameraMatrix=CamMax,
                           distCoeffs=disCef,
                           newCameraMatrix=CamMax
                           )
    # 将矫正后的图片返回出去
    return und_st


if __name__ == '__main__':

    cv2.namedWindow('light_detect_original', cv2.WINDOW_AUTOSIZE)
    # 首先读取摄像头
    # ubuntu下必须要使用cv2.CAP_V4L2这个采集参数，否则会出现1280*720的分辨率
    # cap_camera_wig = cv2.VideoCapture(0, cv2.CAP_V4L2)
    cap_camera_wig = cv2.VideoCapture(0, cv2.CAP_DSHOW)
    # 进入显示视频帧的循环
    while True:

        # 使用read这个函数读取视频帧，  返回bool值和   Mat图片数据
        # ret  //是否读取到数据
        # frame  //读取到的Mat数据
        ret1, frame1 = cap_camera_wig.read()
        # ret2, frame2 = cap_video_wig.read()

        # 这里将获取到的摄像头视频帧进行一个翻转操作(这个操作是后面添加上去的，可以随需求进行调整)
        # frame1 = cv2.flip(frame1, flipCode=1)

        # 如果读取到了，打印日志
        if ret1:

            # 这里进行摄像头的矫正操作
            # und_st = Straighten_the_camera(frame1)
            und_st = frame1
        else:
            und_st = frame1
            print("没有读取到摄像头数据")
        # ----------------------------------------------------------------------------------------------
        #
        # left_img = und_st[:, 0:320]
        #
        # cv2.imshow('left_img', left_img)

        # 显示矫正之后的数据
        cv2.imshow('light_detect_original', und_st)

        # 这里进行红绿灯识别
        # light_detection(und_st)

        # 红绿灯选择识别
        RG_light_choice(und_st)


        # result = light_detection(und_st)
        #
        # if result == 1:
        #     print("Green")
        # elif result == 0:
        #     print("Red")
        # elif result == 2:
        #     print("None")

        # 等待10ms看有没有按键按下，读取指令
        key_wig = cv2.waitKey(1)

        # 这里设置了退出的指令
        if (key_wig & 0xFF) == ord('q'):
            print("Exit the programme")
            break

        # 这里设置了保存图片的指令
        elif (key_wig & 0xFF) == ord('s'):
            cv2.imwrite("/home/nano/wig_rosfile_8_12/cv_test_ws/src/cv_test_pkg/scripts/R_G_light/R_G_light_img.png",
                        und_st,
                        None)
            print("The image of und_st has been saved")

    # 释放VideoCapture这个函数产生的变量 cap_wig
    cap_camera_wig.release()
    # cap_video_wig.release()

    # 释放窗口
    cv2.destroyAllWindows()
